import fasttext

import config


def build_classify_model(by_word=False):
    """模型的创建和保存"""
    # 获取训练数据路径
    data_path = config.classify_corpus_train_by_word_path if by_word else config.classify_corpus_train_path

    # 训练模型
    # wordNgrams是用几个词作为一组，比如wordNgrams=2, ['我', '是', '谁'],训练的时候则将'我', '是'两个词作为一组
    # fasttext.train_supervised创建的模型有predict方法，而train_unsuperivsed创建的模型没有predict方法
    model = fasttext.train_supervised(data_path, wordNgrams=2, epoch=20)

    # 保存模型
    save_model_path = config.classify_model_by_word_path if by_word else config.classify_model_path
    model.save_model(save_model_path)


def get_classify_model(by_word=False):
    """获取模型"""
    # 获取模型路径
    load_model_path = config.classify_model_by_word_path if by_word else config.classify_model_path
    model = fasttext.load_model(load_model_path)

    return model


def evaluate(by_word=False):
    """模型的评估"""
    # 获取测试数据集路径
    test_data_path = config.classify_corpus_test_by_word_path if by_word else config.classify_corpus_test_path

    # 获取加载模型的路径
    load_model_path = config.classify_model_by_word_path if by_word else config.classify_model_path

    # 加载模型
    model = fasttext.load_model(load_model_path)

    # 获取测试数据
    input_data = []
    target = []
    test_data = open(test_data_path, encoding="UTF-8").readlines()
    for line in test_data:
        temp = line.split("__label__")

        input_data.append(temp[0].strip())
        target.append(temp[1].strip())

    # 进行预测
    label, accuracy = model.predict(input_data)
    # print(label)

    # 计算准确率
    pre_true = 0
    for i, j in zip(label, target):
        if i[0].replace("__label__", "").lower() == j.lower():
            pre_true += 1

    # 按词语进行切分：0.99285
    # 按字进行切分：0.99442
    print("accuracy=", pre_true/len(test_data))



